{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "co6fFtLrmEWQ"
      },
      "outputs": [],
      "source": [
        "# Copyright 2022 Google LLC\n",
        "\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License.\n",
        "\n",
        "from colabtools import adhoc_import\n",
        "from typing import Any, Sequence\n",
        "import ml_collections\n",
        "import numpy as np\n",
        "import jax\n",
        "import jax.numpy as jnp\n",
        "import tensorflow as tf\n",
        "import tensorflow_datasets as tfds\n",
        "import functools\n",
        "import itertools\n",
        "from scipy.stats import mode\n",
        "from collections import defaultdict\n",
        "import matplotlib.pyplot as plt\n",
        "from matplotlib.pyplot import cm\n",
        "import time\n",
        "import pickle\n",
        "from sklearn.cluster import KMeans, MiniBatchKMeans\n",
        "from sklearn.decomposition import PCA\n",
        "from sklearn.cluster import AgglomerativeClustering\n",
        "from sklearn.metrics import rand_score, adjusted_rand_score, silhouette_score, adjusted_mutual_info_score\n",
        "from sklearn.mixture import GaussianMixture\n",
        "from clu import preprocess_spec\n",
        "from scipy.special import comb\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8pW1W3WiAt0j"
      },
      "source": [
        "# Compute linear fit of purity vs layers"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vc0DIAySUgzf"
      },
      "outputs": [],
      "source": [
        "from sklearn.linear_model import LinearRegression\n",
        "from sklearn.preprocessing import MinMaxScaler\n",
        "for model_dir in [os.path.join(BASE_DIR, 'breeds/entity13_4_subclasses_400_epochs_ema_0.99_bn_0.99/'),\n",
        "    os.path.join(BASE_DIR, 'breeds/living17_400_epochs_ema_0.99_bn_0.99/'),\n",
        "    os.path.join(BASE_DIR, 'breeds/entity13_8_classes_400_epochs_ema_0.99_bn_0.99/'),\n",
        "    os.path.join(BASE_DIR, 'breeds/living17_8_classes_400_epochs_ema_0.99_bn_0.99/'),\n",
        "    os.path.join(BASE_DIR, 'breeds/breeds_training_level_3_nsubclass_20_400_epochs_ema_0.99/'),\n",
        "    os.path.join(BASE_DIR, 'breeds/nonliving26_400_epochs_ema_0.99_bn_0.99/'),\n",
        "    ]:\n",
        "  all_purity = []\n",
        "  Xs = []\n",
        "  Ys = []\n",
        "  for stage in [1,2,3,4]:\n",
        "    for block in [1,2,3,4,5,6]:\n",
        "      purity_data_path = os.path.join(model_dir, f'class_purity_ckpt_stage{stage}_block{block}.pkl')\n",
        "      if not gfile.Exists(purity_data_path):\n",
        "        continue\n",
        "      with gfile.Open(purity_data_path, 'rb') as f:\n",
        "        all_purity.append(pickle.load(f))\n",
        "  all_purity = np.vstack(all_purity)\n",
        "  for layer_idx in range(all_purity.shape[0]):\n",
        "    Ys.extend(all_purity[layer_idx, :])\n",
        "    Xs.extend([layer_idx+1] * all_purity.shape[1])\n",
        "  Xs = np.array(Xs).reshape(-1, 1)\n",
        "  Ys = np.array(Ys).reshape(-1, 1)\n",
        "  # scaler = MinMaxScaler()\n",
        "  # Xs = scaler.fit_transform(Xs)\n",
        "  # Ys = scaler.fit_transform(Ys)\n",
        "  # print(list(Xs))\n",
        "  # print(list(Ys))\n",
        "  reg = LinearRegression().fit(Xs, Ys.squeeze())\n",
        "  print(model_dir)\n",
        "  print(reg.coef_)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JaJ3sRV4A2iZ"
      },
      "source": [
        "# Plot adjusted mutual info over time"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PFhgJfIVKUf5"
      },
      "outputs": [],
      "source": [
        "for model_dir in [os.path.join(BASE_DIR, 'breeds/living17_400_epochs_ema_0.99_bn_0.99/'), \n",
        "                  os.path.join(BASE_DIR, 'breeds/nonliving26_400_epochs_ema_0.99_bn_0.99/'),\n",
        "                  os.path.join(BASE_DIR, 'breeds/breeds_training_level_3_nsubclass_20_400_epochs_ema_0.99/'),\n",
        "                  os.path.join(BASE_DIR, 'breeds/entity13_4_subclasses_400_epochs_ema_0.99_bn_0.99/'),\n",
        "                  os.path.join(BASE_DIR, 'breeds/entity13_4_subclasses_shuffle_400_epochs_ema_0.99_bn_0.99/')]:\n",
        "  all_ckpt = [f for f in list(gfile.ListDir(model_dir)) if f.startswith('adjusted_mutual_info_') and not f.startswith('adjusted_mutual_info_ckpt_stage')]\n",
        "  all_ckpt = sorted([int(f.split('.')[0].split('_')[-1]) for f in all_ckpt])\n",
        "  #all_ckpt = [c for c in all_ckpt if c % 10 == 1 and c \u003e 1]\n",
        "  ckpt_number = max(all_ckpt)\n",
        "  print(all_ckpt)\n",
        "  n_cols = 6\n",
        "  n_rows = int(np.ceil(len(all_ckpt)/n_cols))\n",
        "  fig, axs = plt.subplots(n_rows, n_cols, figsize=(5*n_cols, 5*n_rows))\n",
        "  print(model_dir)\n",
        "  for idx, ckpt_number in enumerate(all_ckpt):\n",
        "    row = idx // n_cols\n",
        "    col = idx % n_cols\n",
        "    ami_file = os.path.join(model_dir, f'adjusted_mutual_info_ckpt_{ckpt_number}.pkl')\n",
        "    with gfile.Open(ami_file, 'rb') as f:\n",
        "      data = pickle.load(f)\n",
        "    n_subclasses = len(data[1])\n",
        "    fig = plt.figure()\n",
        "    all_subclass_ami = []\n",
        "    color = iter(cm.rainbow(np.linspace(0, 1, n_subclasses)))\n",
        "    for subclass_idx in range(n_subclasses):\n",
        "      subclass_ami = [data[overcluster_factor][subclass_idx] for overcluster_factor in range(1, 6)]\n",
        "      all_subclass_ami.append(subclass_ami)\n",
        "      axs[row, col].plot([1,2,3,4,5], subclass_ami, marker='o', color=next(color))\n",
        "    all_subclass_ami = np.vstack(all_subclass_ami)\n",
        "    axs[row, col].plot([1,2,3,4,5], np.mean(all_subclass_ami, axis=0), marker='o', color='k')\n",
        "    axs[row, col].set_title(f\"Ckpt {ckpt_number}\")#model_dir.split('/')[-2])\n",
        "    axs[row, col].set_xticks([1,2,3,4,5])\n",
        "    axs[row, col].set_xlabel('Overclustering factor')\n",
        "    axs[row, col].set_ylabel('Adjusted mutual information')\n",
        "    #axs[row, col].set_ylim([0.0, 1.0])\n",
        "  plt.show()\n",
        "  plt.clf()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TtxRwZx8BOOa"
      },
      "source": [
        "# Plot adjusted mutual info over layers"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TBB3p0dmoVKx"
      },
      "outputs": [],
      "source": [
        "for model_dir in [os.path.join(BASE_DIR, 'breeds/living17_400_epochs_ema_0.99_bn_0.99/'), \n",
        "                  os.path.join(BASE_DIR, 'breeds/nonliving26_400_epochs_ema_0.99_bn_0.99/'),\n",
        "                  os.path.join(BASE_DIR, 'breeds/breeds_training_level_3_nsubclass_20_400_epochs_ema_0.99/'),\n",
        "                  os.path.join(BASE_DIR, 'breeds/entity13_4_subclasses_400_epochs_ema_0.99_bn_0.99/'),\n",
        "                  os.path.join(BASE_DIR, 'breeds/entity13_4_subclasses_shuffle_400_epochs_ema_0.99_bn_0.99/')]:\n",
        "  print(model_dir)\n",
        "  fig, axs = plt.subplots(4, 6, figsize=(5*6, 5*4))\n",
        "  for stage in [1,2,3,4]:\n",
        "    for block in [1,2,3,4,5,6]:\n",
        "      ami_file = os.path.join(model_dir, f'adjusted_mutual_info_ckpt_stage{stage}_block{block}.pkl')\n",
        "      if not gfile.Exists(ami_file):\n",
        "        continue\n",
        "      with gfile.Open(ami_file, 'rb') as f:\n",
        "        data = pickle.load(f)\n",
        "      n_subclasses = len(data[1])\n",
        "      fig = plt.figure()\n",
        "      all_subclass_ami = []\n",
        "      color = iter(cm.rainbow(np.linspace(0, 1, n_subclasses)))\n",
        "      for subclass_idx in range(n_subclasses):\n",
        "        subclass_ami = [data[overcluster_factor][subclass_idx] for overcluster_factor in range(1, 6)]\n",
        "        all_subclass_ami.append(subclass_ami)\n",
        "        axs[stage-1, block-1].plot([1,2,3,4,5], subclass_ami, marker='o', c=next(color))\n",
        "      all_subclass_ami = np.vstack(all_subclass_ami)\n",
        "      axs[stage-1, block-1].plot([1,2,3,4,5], np.mean(all_subclass_ami, axis=0), marker='o', color='k', linewidth=2)\n",
        "      axs[stage-1, block-1].set_xticks([1,2,3,4,5])\n",
        "      axs[stage-1, block-1].set_xlabel('Overclustering factor')\n",
        "      axs[stage-1, block-1].set_ylabel('Adjusted mutual information')\n",
        "      axs[stage-1, block-1].set_title(f\"stage{stage}_block{block}\")\n",
        "      #axs[stage-1, block-1].set_ylim([0.0, 1.0])\n",
        "  plt.show()\n",
        "  plt.clf()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5TQ68ctiNqLZ"
      },
      "source": [
        "# Plot purity vs training steps, and adjust_rand_idex of clusters from successive checkpoints vs steps"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1sJCOhCHgkzw"
      },
      "outputs": [],
      "source": [
        "ret = make_breeds_dataset(\"living17\", BREEDS_INFO_DIR, BREEDS_INFO_DIR,\n",
        "                    split=None,\n",
        "                    #num_subclasses=4, shuffle_subclasses=True\n",
        "                    )\n",
        "superclasses, subclass_split, label_map = ret\n",
        "train_subclasses = subclass_split[0]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mP3NxqyugnPe"
      },
      "outputs": [],
      "source": [
        "all_purity = []\n",
        "all_ckpt_numbers = list(range(3, 42, 2))\n",
        "Xs = []\n",
        "for ckpt_number in all_ckpt_numbers:\n",
        "  purity_data_path = fos.path.join(BASE_DIR, 'breeds/1/class_purity_ckpt_{ckpt_number}.pkl'\n",
        "  if not gfile.Exists(purity_data_path):\n",
        "    continue\n",
        "  Xs.append(ckpt_number)\n",
        "  with gfile.Open(purity_data_path, 'rb') as f:\n",
        "    all_purity.append(pickle.load(f))\n",
        "all_purity = np.vstack(all_purity)\n",
        "print(all_purity.shape)\n",
        "\n",
        "color = iter(cm.rainbow(np.linspace(0, 1, all_purity.shape[1])))\n",
        "fig = plt.figure(figsize=(15,12))\n",
        "for i in range(all_purity.shape[1]):\n",
        "  c = next(color)\n",
        "  plt.plot([100*(c-1) for c in Xs], all_purity[:, i], label=label_map[i], color=c, marker='o')\n",
        "\n",
        "plt.ylim((0, 1.0))\n",
        "plt.xlabel(\"Step\")\n",
        "plt.ylabel(\"Cluster purity\")\n",
        "plt.legend()\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WbJ1Kfqgh67R"
      },
      "outputs": [],
      "source": [
        "all_rand_index = [[] for i in range(len(train_subclasses))]\n",
        "Xs = list(range(5, 42, 2))\n",
        "for ckpt_number in Xs:\n",
        "  with gfile.Open(fos.path.join(BASE_DIR, 'breeds/living17_400_epochs_ema_0.99_bn_0.99_squared_loss/clf_labels_ckpt_{ckpt_number}.pkl'), 'rb') as f:\n",
        "    all_clf_labels = pickle.load(f)\n",
        "  with gfile.Open(fos.path.join(BASE_DIR, 'breeds/living17_400_epochs_ema_0.99_bn_0.99_squared_loss/clf_labels_ckpt_{ckpt_number-2}.pkl'), 'rb') as f:\n",
        "    all_clf_labels_prev = pickle.load(f)\n",
        "  for i in range(len(train_subclasses)):\n",
        "    all_rand_index[i].append(adjusted_rand_score(all_clf_labels[i], all_clf_labels_prev[i]))\n",
        "  \n",
        "color = iter(cm.rainbow(np.linspace(0, 1, all_purity.shape[1])))\n",
        "fig = plt.figure(figsize=(15,12))\n",
        "for i in range(len(train_subclasses)):\n",
        "  c = next(color)\n",
        "  plt.plot([100*(c-1) for c in Xs], all_rand_index[i], label=label_map[i], color=c, marker='o')\n",
        "\n",
        "plt.xlabel(\"Step\", fontsize=16)\n",
        "plt.ylabel(\"Adjusted Rand Index\", fontsize=16)\n",
        "plt.legend()\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bbkB4B0wBh7h"
      },
      "source": [
        "# Plot purity vs layers, and adjust_rand_idex of clusters from successive layers vs steps"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "773bTtObg84k"
      },
      "outputs": [],
      "source": [
        "all_purity = []\n",
        "Xs = []\n",
        "X_ticks = []\n",
        "for stage in [1,2,3,4]:\n",
        "  for block in [1,2,3,4,5,6]:\n",
        "    purity_data_path = fos.path.join(BASE_DIR, 'breeds/1/class_purity_ckpt_stage{stage}_block{block}.pkl')\n",
        "    if not gfile.Exists(purity_data_path):\n",
        "      continue\n",
        "    X_ticks.append(f\"stage{stage}_block{block}\")\n",
        "    with gfile.Open(purity_data_path, 'rb') as f:\n",
        "      all_purity.append(pickle.load(f))\n",
        "all_purity = np.vstack(all_purity)\n",
        "print(all_purity.shape)\n",
        "\n",
        "colors = iter(cm.rainbow(np.linspace(0, 1, all_purity.shape[1])))\n",
        "fig = plt.figure(figsize=(15,12))\n",
        "for i in range(all_purity.shape[1]):\n",
        "  color = next(colors)\n",
        "  plt.plot(range(len(X_ticks)), all_purity[:, i], label=label_map[i], color=color, marker='o')\n",
        "\n",
        "# Xs = list(range(16))\n",
        "# Ys = [x * 0.01190781 + 0.5176538461538461 for x in Xs]\n",
        "# plt.plot(Xs, Ys, color='k')\n",
        "plt.xticks(ticks=range(len(X_ticks)), labels=X_ticks, rotation=90)\n",
        "plt.ylim((0, 1.0))\n",
        "plt.xlabel(\"Layer\")\n",
        "plt.ylabel(\"Cluster purity\")\n",
        "plt.legend()\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AfvAUXUzhkW7"
      },
      "outputs": [],
      "source": [
        "all_rand_index = [[] for i in range(len(train_subclasses))]\n",
        "for i in range(1, len(X_ticks)):\n",
        "  current_block = X_ticks[i]\n",
        "  prev_block = X_ticks[i-1]\n",
        "  with gfile.Open(fos.path.join(BASE_DIR, 'breeds/1/clf_labels_ckpt_{current_block}.pkl'), 'rb') as f:\n",
        "    all_clf_labels = pickle.load(f)\n",
        "  with gfile.Open(fos.path.join(BASE_DIR, 'breeds/1/clf_labels_ckpt_{prev_block}.pkl'), 'rb') as f:\n",
        "    all_clf_labels_prev = pickle.load(f)\n",
        "  for i in range(len(train_subclasses)):\n",
        "    all_rand_index[i].append(adjusted_rand_score(all_clf_labels[i], all_clf_labels_prev[i]))\n",
        "  \n",
        "colors = iter(cm.rainbow(np.linspace(0, 1, all_purity.shape[1])))\n",
        "fig = plt.figure(figsize=(15,12))\n",
        "for i in range(len(train_subclasses)):\n",
        "  color = next(colors)\n",
        "  plt.plot(range(1, len(X_ticks)), all_rand_index[i], label=label_map[i], color=color, marker='o')\n",
        "\n",
        "plt.xticks(ticks=range(1, len(X_ticks)), labels=X_ticks[1:], rotation=90)\n",
        "plt.xlabel(\"Layer\", fontsize=16)\n",
        "plt.ylabel(\"Adjusted Rand Index\", fontsize=16)\n",
        "plt.legend()\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NJoo0uUS5-wl"
      },
      "source": [
        "# Load eval data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_bTCxV_vEpzV"
      },
      "outputs": [],
      "source": [
        "ret = make_breeds_dataset(\"living17\", BREEDS_INFO_DIR, BREEDS_INFO_DIR,\n",
        "                          split=None)\n",
        "superclasses, subclass_split, label_map = ret\n",
        "train_subclasses = subclass_split[0]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AFDKat-K9kx8"
      },
      "outputs": [],
      "source": [
        "def load_dataset(dataset_name, train_subclasses=None):\n",
        "  if dataset_name == 'imagenet':\n",
        "    all_subclasses = list(itertools.chain(*train_subclasses))\n",
        "    new_label_map = {}\n",
        "    for subclass_idx, sub in enumerate(all_subclasses):\n",
        "      new_label_map.update({sub: subclass_idx})\n",
        "    print(new_label_map)\n",
        "    lookup_table = tf.lookup.StaticHashTable(\n",
        "        initializer=tf.lookup.KeyValueTensorInitializer(\n",
        "            keys=tf.constant(list(new_label_map.keys()), dtype=tf.int64),\n",
        "            values=tf.constant(list(new_label_map.values()), dtype=tf.int64),\n",
        "        ),\n",
        "        default_value=tf.constant(-1, dtype=tf.int64))\n",
        "\n",
        "    dataset_builder = tfds.builder(\"imagenet2012\", try_gcs=True)\n",
        "    eval_preprocess = preprocess_spec.PreprocessFn([\n",
        "        RescaleValues(),\n",
        "        ResizeSmall(256),\n",
        "        CentralCrop(224),\n",
        "        ], only_jax_types=True)\n",
        "    num_classes = len(train_subclasses)\n",
        "  else:\n",
        "    dataset_builder = tfds.builder(\"celeb_a\", try_gcs=True)\n",
        "    eval_preprocess = preprocess_spec.PreprocessFn([\n",
        "        RescaleValues(),\n",
        "        ResizeSmall(256),\n",
        "        CentralCrop(224),\n",
        "        LabelMapping()\n",
        "        ], only_jax_types=True)\n",
        "    num_classes = 2\n",
        "\n",
        "  dataset_options = tf.data.Options()\n",
        "  dataset_options.experimental_optimization.map_parallelization = True\n",
        "  dataset_options.experimental_threading.private_threadpool_size = 48\n",
        "  dataset_options.experimental_threading.max_intra_op_parallelism = 1\n",
        "\n",
        "  read_config = tfds.ReadConfig(shuffle_seed=None, options=dataset_options)\n",
        "  eval_ds = dataset_builder.as_dataset(\n",
        "      split=tfds.Split.VALIDATION,\n",
        "      shuffle_files=False,\n",
        "      read_config=read_config,\n",
        "      decoders=None)\n",
        "\n",
        "  batch_size = 64\n",
        "  if dataset_name == 'imagenet':\n",
        "    eval_ds = eval_ds.filter(functools.partial(predicate, all_subclasses=all_subclasses))\n",
        "  eval_ds = eval_ds.cache()\n",
        "  eval_ds = eval_ds.map(eval_preprocess, num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
        "  eval_ds = eval_ds.batch(batch_size, drop_remainder=False)\n",
        "  eval_ds = eval_ds.prefetch(tf.data.experimental.AUTOTUNE)\n",
        "  return eval_ds, num_classes"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "x_khNgu_DoAP"
      },
      "outputs": [],
      "source": [
        "DATASET = 'imagenet'#'celeb_a'\n",
        "eval_ds, num_classes = load_dataset(DATASET, train_subclasses=train_subclasses)\n",
        "if DATASET == 'celeb_a':\n",
        "  train_subclasses = [[0], [1]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kdsxp5e48si8"
      },
      "source": [
        "# Clustering algorithms"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0voo_LcJEfDL"
      },
      "outputs": [],
      "source": [
        "def get_learning_rate(step: int,\n",
        "                      *,\n",
        "                      base_learning_rate: float,\n",
        "                      steps_per_epoch: int,\n",
        "                      num_epochs: int,\n",
        "                      warmup_epochs: int = 5):\n",
        "  \"\"\"Cosine learning rate schedule.\"\"\"\n",
        "  logging.info(\n",
        "      \"get_learning_rate(step=%s, base_learning_rate=%s, steps_per_epoch=%s, num_epochs=%s\",\n",
        "      step, base_learning_rate, steps_per_epoch, num_epochs)\n",
        "  if steps_per_epoch \u003c= 0:\n",
        "    raise ValueError(f\"steps_per_epoch should be a positive integer but was \"\n",
        "                     f\"{steps_per_epoch}.\")\n",
        "  if warmup_epochs \u003e= num_epochs:\n",
        "    raise ValueError(f\"warmup_epochs should be smaller than num_epochs. \"\n",
        "                     f\"Currently warmup_epochs is {warmup_epochs}, \"\n",
        "                     f\"and num_epochs is {num_epochs}.\")\n",
        "  epoch = step / steps_per_epoch\n",
        "  lr = cosine_decay(base_learning_rate, epoch - warmup_epochs,\n",
        "                    num_epochs - warmup_epochs)\n",
        "  warmup = jnp.minimum(1., epoch / warmup_epochs)\n",
        "  return lr * warmup\n",
        "\n",
        "config = get_config()\n",
        "learning_rate_fn = functools.partial(\n",
        "      get_learning_rate,\n",
        "      base_learning_rate=0.1,\n",
        "      steps_per_epoch=40,\n",
        "      num_epochs=config.num_epochs,\n",
        "      warmup_epochs=config.warmup_epochs)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tfl1KvfMEBm1"
      },
      "outputs": [],
      "source": [
        "model_dir = os.path.join(BASE_DIR, 'breeds/living17_400_epochs_ema_0.99_bn_0.99_squared_loss_hyperparam_tuned/')\n",
        "checkpoint_path = os.path.join(model_dir, 'checkpoints-0/ckpt-41.flax')\n",
        "model, state = create_train_state(config, jax.random.PRNGKey(0), input_shape=(8, 224, 224, 3), num_classes=num_classes, learning_rate_fn=learning_rate_fn)\n",
        "state = checkpoints.restore_checkpoint(checkpoint_path, state)\n",
        "print(\"step:\", state.step)\n",
        "\n",
        "all_intermediates = []\n",
        "all_subclass_labels = []\n",
        "all_images = []\n",
        "for step, batch in enumerate(eval_ds):\n",
        "  if step % 20 == 0:\n",
        "    print(step)\n",
        "  intermediates = predict(model, state, batch)\n",
        "  labels = batch['label'].numpy()\n",
        "  bs = labels.shape[0]\n",
        "  all_subclass_labels.append(labels)\n",
        "  all_images.append(batch['image'].numpy())\n",
        "  all_intermediates.append(np.mean(intermediates['stage4']['__call__'][0], axis=(1,2)).reshape(bs, -1))\n",
        "\n",
        "all_intermediates = np.vstack(all_intermediates)\n",
        "all_subclass_labels = np.hstack(all_subclass_labels)\n",
        "all_images = np.vstack(all_images)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cDNYGa648vTm"
      },
      "outputs": [],
      "source": [
        "cluster_each_superclass = True\n",
        "#overcluster_factor, n_subclasses = 5, 20\n",
        "if isinstance(all_intermediates, list):\n",
        "  all_intermediates = np.vstack(all_intermediates)\n",
        "if isinstance(all_subclass_labels, list):\n",
        "  all_subclass_labels = np.hstack(all_subclass_labels)\n",
        "if isinstance(all_images, list):\n",
        "  all_images = np.vstack(all_images)\n",
        "if len(all_filenames) \u003e 0:\n",
        "  all_filenames = np.hstack(all_filenames)\n",
        "  all_filenames = [f.decode(\"utf-8\") for f in all_filenames]\n",
        "print(all_intermediates.shape)\n",
        "print(all_subclass_labels.shape)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3_-6St7LTUTK"
      },
      "source": [
        "### Kmeans (with PCA 50 dim)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CshLL8HsMzAu"
      },
      "outputs": [],
      "source": [
        "all_intermediates_normalized = all_intermediates - np.mean(all_intermediates, axis=0)\n",
        "pca = PCA(n_components=50)\n",
        "all_intermediates_pca = pca.fit_transform(all_intermediates_normalized)\n",
        "\n",
        "for overcluster_factor in [0.1, 0.25, 0.5, 1, 3, 5, 10]:\n",
        "  all_clfs = []\n",
        "  all_sil_scores = []\n",
        "  if cluster_each_superclass:\n",
        "    for subclasses in train_subclasses:-\n",
        "      subclass_idx = np.array([i for i in range(len(all_subclass_labels)) if all_subclass_labels[i] in subclasses])\n",
        "      kmeans = KMeans(n_clusters=int(len(subclasses)*overcluster_factor), random_state=1).fit(all_intermediates_pca[subclass_idx])\n",
        "      all_sil_scores.append(silhouette_score(all_intermediates_pca[subclass_idx], kmeans.labels_))\n",
        "      all_clfs.append(kmeans)\n",
        "    print(all_sil_scores)\n",
        "    print(f\"overcluster factor = {overcluster_factor}, mean sil score = {np.mean(all_sil_scores)}\")\n",
        "  else:\n",
        "    kmeans = KMeans(n_clusters=int(len(all_subclasses)*overcluster_factor), random_state=1).fit(all_intermediates_pca)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zK3PW1PJTZRF"
      },
      "source": [
        "### Minibatch Kmeans"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Qe4ZjL0euPZg"
      },
      "outputs": [],
      "source": [
        "start = time.time()\n",
        "kmeans_minibatch = MiniBatchKMeans(n_clusters=len(all_subclasses)*overcluster_factor, random_state=0, batch_size=1024, max_iter=300, max_no_improvement=100).fit(all_intermediates)\n",
        "end = time.time()\n",
        "print(\"Time taken:\", end - start)\n",
        "print(\"number of clusters found:\", len(set(kmeans_minibatch.labels_)))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IOeIkVHkTb8M"
      },
      "source": [
        "### Hierarchical clustering"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tKI1-89Ek7Qq"
      },
      "outputs": [],
      "source": [
        "for overcluster_factor in [0.1, 0.25, 0.5, 1, 3, 5, 10]:\n",
        "  all_clfs = []\n",
        "  all_sil_scores = []\n",
        "  if cluster_each_superclass:\n",
        "    for subclasses in train_subclasses:\n",
        "      subclass_idx = np.array([i for i in range(len(all_subclass_labels)) if all_subclass_labels[i] in subclasses])\n",
        "      hier_clustering = AgglomerativeClustering(n_clusters=int(len(subclasses)*overcluster_factor),\n",
        "                                          linkage='ward').fit(all_intermediates[subclass_idx])\n",
        "      all_sil_scores.append(silhouette_score(all_intermediates[subclass_idx], hier_clustering.labels_))\n",
        "      all_clfs.append(hier_clustering)\n",
        "    print(f\"overcluster factor = {overcluster_factor}, mean sil score = {np.mean(all_sil_scores)}\")\n",
        "  else:\n",
        "    hier_clustering = AgglomerativeClustering(n_clusters=int(len(all_subclasses)*overcluster_factor),\n",
        "                                          linkage='ward').fit(all_intermediates)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qpXNXrIFhr4V"
      },
      "source": [
        "### GMM"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "E5W6MeaYhxuH"
      },
      "outputs": [],
      "source": [
        "for overcluster_factor in [3, 5, 10, 15, 20, 25]:\n",
        "  all_clfs = []\n",
        "  all_bic_scores = []\n",
        "  if cluster_each_superclass:\n",
        "    for subclasses in train_subclasses:\n",
        "      subclass_idx = np.array([i for i in range(len(all_subclass_labels)) if all_subclass_labels[i] in subclasses])\n",
        "      gmm = GaussianMixture(n_components=len(subclasses)*overcluster_factor, covariance_type='tied', \n",
        "                            random_state=1, reg_covar=1e-5)\n",
        "      gmm.fit(all_intermediates[subclass_idx])\n",
        "      all_clfs.append(gmm.predict(all_intermediates[subclass_idx]))\n",
        "      all_bic_scores.append(gmm.bic(all_intermediates[subclass_idx]))\n",
        "    print(f\"overcluster factor = {overcluster_factor}, mean bic score = {np.mean(all_bic_scores)}\")\n",
        "  else:\n",
        "    gmm = GaussianMixture(n_components=len(subclasses)*overcluster_factor, covariance_type='full', random_state=1)\n",
        "    gmm.fit(all_intermediates)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lTR6VhtH8v4q"
      },
      "source": [
        "# Analyze images and labels in each cluster"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fLUGzxmC-_wv"
      },
      "outputs": [],
      "source": [
        "def show_images_horizontally(images, labels, super_label=None, max_samples=20):\n",
        "  if max_samples is not None:\n",
        "    number_of_files = min(images.shape[0], max_samples)\n",
        "    fig = plt.figure(figsize=(1.5 * number_of_files, 1.5))\n",
        "    n_rows = 1\n",
        "    n_cols = number_of_files\n",
        "  else:\n",
        "    number_of_files = images.shape[0]\n",
        "    n_cols = 10\n",
        "    n_rows = int(np.ceil(number_of_files/n_cols))\n",
        "    fig = plt.figure(figsize=(1.5 * n_cols, 1.5 * n_rows))\n",
        "                          \n",
        "  for i in range(number_of_files):\n",
        "    axes = fig.add_subplot(n_rows, n_cols, i+1)\n",
        "    axes.imshow(images[i])\n",
        "    if isinstance(labels, list):\n",
        "      axes.set_title(labels[i])\n",
        "    else:\n",
        "      axes.set_title(labels)\n",
        "  if super_label:\n",
        "    plt.title(super_label)\n",
        "  plt.show()\n",
        "\n",
        "def visualize_all_clusters(all_images, cluster_labels, all_subclass_labels, hier, gender_labels):\n",
        "  for label in list(set(cluster_labels)):\n",
        "    instance_idx = np.where(cluster_labels == label)[0]\n",
        "    subclass_labels = all_subclass_labels[instance_idx]\n",
        "    mode_stats = mode(subclass_labels)\n",
        "    # if mode_stats[1][0] == 1:\n",
        "    #   # Skip clusters with no dominant subclass\n",
        "    #   continue\n",
        "    cluster_images = all_images[instance_idx]\n",
        "    cluster_label = mode_stats[0][0]\n",
        "    if hier:\n",
        "      cluster_label_str = hier.LEAF_NUM_TO_NAME[cluster_label].split(',')[0]\n",
        "      # if cluster_label not in subclasses_to_visualize:\n",
        "      #   # only visualize clusters where the dominant subclass belongs to subclasses_to_visualize\n",
        "      #   continue\n",
        "      subclass_labels = [hier.LEAF_NUM_TO_NAME[s].split(',')[0] for s in subclass_labels]\n",
        "      show_images_horizontally(cluster_images, subclass_labels, super_label=cluster_label_str, max_samples=10)\n",
        "    elif gender_labels:\n",
        "      male_pc = sum(gender_labels[instance_idx]) * 100.0 / len(instance_idx)\n",
        "      cluster_label_str = f'Male ratio = {male_pc}%'\n",
        "      subclass_labels = ['' for s in subclass_labels]\n",
        "      show_images_horizontally(cluster_images, subclass_labels, super_label=cluster_label_str, max_samples=10)\n",
        "    else:\n",
        "      show_images_horizontally(cluster_images, subclass_labels, super_label='', max_samples=10)\n",
        "\n",
        "\n",
        "def visualize_cluster(cluster_idx, clusters, all_images, all_subclass_labels, hier):\n",
        "  instance_idx = np.where(clusters == cluster_idx)[0]\n",
        "  cluster_images = all_images[instance_idx]\n",
        "  subclass_labels = all_subclass_labels[instance_idx]\n",
        "  mode_stats = mode(subclass_labels)\n",
        "  if mode_stats[1][0] == 1:\n",
        "    print(\"No dominant subclass\")\n",
        "  else:\n",
        "    cluster_label = mode_stats[0][0]\n",
        "    cluster_label_str = hier.LEAF_NUM_TO_NAME[cluster_label].split(',')[0]\n",
        "    print(\"Dominant subclass:\", cluster_label_str)\n",
        "  subclass_labels = [hier.LEAF_NUM_TO_NAME[s].split(',')[0] for s in subclass_labels]\n",
        "  show_images_horizontally(cluster_images, subclass_labels, max_samples=None)\n",
        "\n",
        "def visualize_error(cluster_label_str, pred_str, clusters, all_subclasses, all_images, hier):\n",
        "  for s in all_subclasses:\n",
        "    if hier.LEAF_NUM_TO_NAME[s].startswith(cluster_label_str):\n",
        "      cluster_label = s\n",
        "  for s in all_subclasses:\n",
        "    if hier.LEAF_NUM_TO_NAME[s].startswith(pred_str):\n",
        "      pred = s\n",
        "  for cluster_idx in list(set(clusters)):\n",
        "    instance_idx = np.where(clusters == cluster_idx)[0]\n",
        "    subclass_labels = all_subclass_labels[instance_idx]\n",
        "    mode_stats = mode(subclass_labels)\n",
        "    if mode_stats[1][0] == 1:\n",
        "      continue\n",
        "    if mode_stats[0][0] != cluster_label:\n",
        "      continue\n",
        "    if pred not in subclass_labels:\n",
        "      continue\n",
        "    cluster_images = all_images[instance_idx]\n",
        "    subclass_labels = [hier.LEAF_NUM_TO_NAME[s].split(',')[0] for s in subclass_labels]\n",
        "    show_images_horizontally(cluster_images, subclass_labels, max_samples=None)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cTHDz13YkHlC"
      },
      "source": [
        "### Save image filenames from each cluster"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MjguWAvvg8IF"
      },
      "outputs": [],
      "source": [
        "cluster_filenames_dict = {}\n",
        "for cluster_idx in set(clf.labels_):\n",
        "  instance_idx = np.where(clf.labels_ == cluster_idx)[0]\n",
        "  filenames = [all_filenames[i] for i in instance_idx]\n",
        "  cluster_filenames_dict[cluster_idx] = filenames\n",
        "cluster_filenames_dict['ckpt'] = checkpoint_path\n",
        "\n",
        "with gfile.Open(os.path.join(BASE_DIR, 'breeds/cluster_filenames'), 'wb') as f:\n",
        "  pickle.dump(cluster_filenames_dict, f)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VMWz2FCQkR2U"
      },
      "source": [
        "### Compute some clustering metrics"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "g-eAGdVbgtbQ"
      },
      "outputs": [],
      "source": [
        "def compute_agg_metrics(clusters, classes, hier, all_subclasses):\n",
        "  tp_plus_fp = comb(np.bincount(clusters), 2).sum()\n",
        "  tp_plus_fn = comb(np.bincount(classes), 2).sum()\n",
        "  A = np.c_[(clusters, classes)]\n",
        "  tp = sum(comb(np.bincount(A[A[:, 0] == i, 1]), 2).sum()\n",
        "            for i in set(clusters))\n",
        "  fp = tp_plus_fp - tp\n",
        "  fn = tp_plus_fn - tp\n",
        "  tn = comb(len(A), 2) - tp - fp - fn\n",
        "  print(\"TP=\", tp)\n",
        "  print(\"TN=\", tn)\n",
        "  print(\"FP=\", fp)\n",
        "  print(\"FN=\", fn)\n",
        "  rand_index_score = (tp + tn) / (tp + fp + fn + tn)\n",
        "  f_score = tp / (tp + 0.5 * (fp + fn))\n",
        "  beta = 0.5 # downweighs fn\n",
        "  f_beta = (1 + beta**2) * tp / ((1 + beta**2) * tp + beta**2 * fn + fp)\n",
        "\n",
        "  all_cluster_sizes = []\n",
        "  all_cluster_label_count = defaultdict(int)\n",
        "  clusters_with_no_mode = []\n",
        "  n_cluster_points = 0\n",
        "  confusion_matrix = np.zeros((len(all_subclasses), len(all_subclasses)))\n",
        "  for cluster_idx in set(clusters):\n",
        "    instance_idx = np.where(clusters == cluster_idx)[0]\n",
        "    all_cluster_sizes.append(len(instance_idx))\n",
        "    subclass_labels = classes[instance_idx]\n",
        "    mode_stats = mode(subclass_labels)\n",
        "    n_cluster_points += mode_stats[1][0]\n",
        "    assert(mode_stats[0][0] in all_subclasses)\n",
        "    if mode_stats[1][0] == 1: # no single class dominating\n",
        "      if len(instance_idx) == 1:\n",
        "        all_cluster_label_count[mode_stats[0][0]] += 1\n",
        "      else:\n",
        "        clusters_with_no_mode.append(cluster_idx)\n",
        "        continue\n",
        "    else:\n",
        "      all_cluster_label_count[mode_stats[0][0]] += 1\n",
        "    for l in subclass_labels:\n",
        "      if l != mode_stats[0][0]:\n",
        "        i = all_subclasses.index(l)\n",
        "        j = all_subclasses.index(mode_stats[0][0])\n",
        "        confusion_matrix[j][i] += 1 #row=true label, col = wrong pred\n",
        "\n",
        "  if cluster_each_superclass:\n",
        "    percent_subclass_covered = len(all_cluster_label_count.keys()) / len(all_subclasses)\n",
        "    print(\"Proportion of subclasses covered:\", percent_subclass_covered)\n",
        "    if percent_subclass_covered \u003c 1:\n",
        "      subclasses_uncovered = set(all_subclasses) - set(all_cluster_label_count.keys())\n",
        "      print(\"Classes that don't dominate any cluster\", [hier.LEAF_NUM_TO_NAME[s] for s in list(subclasses_uncovered)])\n",
        "  print(f\"{len(clusters_with_no_mode)} clusters with no mode:\")\n",
        "  print(clusters_with_no_mode)\n",
        "  purity = n_cluster_points / len(clusters)\n",
        "\n",
        "  # plt.hist(all_cluster_sizes, bins=list(range(max(all_cluster_sizes)+1)))\n",
        "  # plt.xlabel(\"cluster size\")\n",
        "  # plt.ylabel(\"number of clusters\")\n",
        "  # plt.show()\n",
        "  return purity, rand_index_score, f_score, f_beta, confusion_matrix\n",
        "\n",
        "\n",
        "all_confusion_matrix = []\n",
        "purity_avg, rand_index_score_avg, f_score_avg, f_beta_avg = 0, 0, 0, 0\n",
        "for i, clf in enumerate(all_clfs):\n",
        "  if cluster_each_superclass:\n",
        "    print(\"----------------------------SUPERCLASS:\", label_map[i])\n",
        "    subclasses = train_subclasses[i]\n",
        "    subclass_idx = np.array([i for i in range(len(all_subclass_labels)) if all_subclass_labels[i] in subclasses])\n",
        "    subclass_labels = all_subclass_labels[subclass_idx]\n",
        "  else:\n",
        "    subclasses = all_subclasses\n",
        "    subclass_labels = all_subclass_labels\n",
        "  purity, rand_index_score, f_score, f_beta, confusion_matrix = compute_agg_metrics(clf.labels_, \n",
        "                                                                                  subclass_labels, hier,\n",
        "                                                                                  subclasses)\n",
        "  all_confusion_matrix.append(confusion_matrix)\n",
        "  print(\"Purity:\", purity)\n",
        "  print(\"Rand index score:\", rand_index_score)\n",
        "  print(\"F1 score:\", f_score)\n",
        "  print(\"F_beta score:\", f_beta)\n",
        "  purity_avg += purity\n",
        "  rand_index_score_avg += rand_index_score\n",
        "  f_score_avg += f_score\n",
        "  f_beta_avg += f_beta\n",
        "\n",
        "purity_avg /= len(all_clfs)\n",
        "rand_index_score_avg /= len(all_clfs)\n",
        "f_score_avg /= len(all_clfs)\n",
        "f_beta_avg /= len(all_clfs)\n",
        "print(\"----------------------------AVERAGE\")\n",
        "print(\"Purity:\", purity_avg)\n",
        "print(\"Rand index score:\", rand_index_score_avg)\n",
        "print(\"F1 score:\", f_score_avg)\n",
        "print(\"F_beta score:\", f_beta_avg)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CS00ohA8_S3y"
      },
      "source": [
        "### Visualize subclasses that tend to be confused together"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "amhEXXJI9LNm"
      },
      "outputs": [],
      "source": [
        "visualize_error('cucumber', 'bell pepper', clf.labels_, all_subclasses, all_images, hier)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "X3ptfy2TAVz6"
      },
      "outputs": [],
      "source": [
        "visualize_error('orange', 'lemon', clf.labels_, all_subclasses, all_images, hier)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MOfAoXw2pG18"
      },
      "outputs": [],
      "source": [
        "visualize_error('Samoyed', 'Persian cat', clf.labels_, all_subclasses, all_images, hier)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vJOgTUx6pp90"
      },
      "source": [
        "### Visualize clusters with no dominant subclass\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LxMFjj8CoqUZ"
      },
      "outputs": [],
      "source": [
        "visualize_cluster(24, clf.labels_, all_images, all_subclass_labels, hier)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xi5dAlT8RMji"
      },
      "source": [
        "### Visualize all clusters"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zY6ci7Uek5Cv"
      },
      "outputs": [],
      "source": [
        "ret = make_breeds_dataset(\"living17\", BREEDS_INFO_DIR, BREEDS_INFO_DIR,\n",
        "                         split=None)\n",
        "superclasses, subclass_split, label_map = ret\n",
        "train_subclasses = subclass_split[0]\n",
        "hier = ClassHierarchy()\n",
        "eval_ds, num_classes = load_dataset('imagenet', train_subclasses=train_subclasses)\n",
        "all_subclass_labels = []\n",
        "all_images = []\n",
        "for step, batch in enumerate(eval_ds):\n",
        "  if step % 20 == 0:\n",
        "    print(step)\n",
        "  labels = batch['label'].numpy()\n",
        "  bs = labels.shape[0]\n",
        "  all_subclass_labels.append(labels)\n",
        "  all_images.append(batch['image'].numpy())\n",
        "\n",
        "all_subclass_labels = np.hstack(all_subclass_labels)\n",
        "all_labels = all_subclass_labels\n",
        "all_images = np.vstack(all_images)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "11CSMXZXjQ4-"
      },
      "outputs": [],
      "source": [
        "model_path = os.path.join(BASE_DIR, 'breeds/living17_400_epochs_ema_0.99_bn_0.99/')\n",
        "with gfile.Open(os.path.join(model_path, 'clf_labels_ckpt_171_ocfactor_3.pkl'), 'rb') as f:\n",
        "  all_clf_labels = pickle.load(f)\n",
        "superclass_idx = 15\n",
        "print(\"----------SUPERCLASS:\", label_map[superclass_idx])\n",
        "subclass_idx = np.array([i for i in range(len(all_subclass_labels)) if all_subclass_labels[i] in train_subclasses[superclass_idx]])\n",
        "visualize_all_clusters(all_images[subclass_idx], all_clf_labels[superclass_idx], all_subclass_labels[subclass_idx], hier, gender_labels=None)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SB6Ybko1ndDt"
      },
      "outputs": [],
      "source": [
        "if DATASET == 'celeb_a':\n",
        "  all_gender_labels = []\n",
        "  for step, batch in enumerate(eval_ds):\n",
        "    gender = batch['attributes']['Male'].numpy().astype(int)\n",
        "    all_gender_labels.append(gender)\n",
        "  all_gender_labels = np.hstack(all_gender_labels)\n",
        "  print(all_gender_labels.shape)\n",
        "\n",
        "model_path = os.path.join(BASE_DIR, 'celebA/lr_0.0001_reg_0.0001_ema_0.99/')\n",
        "with gfile.Open(os.path.join(model_path, 'clf_labels_ckpt_20.pkl'), 'rb') as f:\n",
        "  all_clf_labels = pickle.load(f)\n",
        "all_clf_labels = [clf.labels_ for clf in all_clfs]\n",
        "class_idx = 0\n",
        "subclass_idx = np.array([i for i in range(len(all_labels)) if all_labels[i] == class_idx])\n",
        "print(f\"---------Male ratio in this class: {sum(all_gender_labels[subclass_idx]) * 100.0/ len(subclass_idx)}%\")\n",
        "visualize_all_clusters(all_images[subclass_idx], all_clf_labels[class_idx], all_subclass_labels[subclass_idx], hier=None, gender_labels=all_gender_labels[subclass_idx])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3CaVOZ9qCcx0"
      },
      "source": [
        "## Visualize cluster reassignment"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cxTHE_o69aJe"
      },
      "outputs": [],
      "source": [
        "superclass_idx = 0\n",
        "print(\"----------SUPERCLASS:\", label_map[superclass_idx])\n",
        "subclass_idx = np.array([i for i in range(len(all_subclass_labels)) if all_subclass_labels[i] in train_subclasses[superclass_idx]])\n",
        "subclass_labels = all_subclass_labels[subclass_idx]\n",
        "n_data = len(subclass_idx)\n",
        "\n",
        "model_path = os.path.join(BASE_DIR, 'breeds/breeds_training_level_3_nsubclass_20_400_epochs_ema_0.99/')\n",
        "with gfile.Open(os.path.join(model_path, 'clf_labels_ckpt_151.pkl'), 'rb') as f:\n",
        "  all_clf_labels_prev = pickle.load(f)[superclass_idx]\n",
        "with gfile.Open(os.path.join(model_path, 'clf_labels_ckpt_161.pkl'), 'rb') as f:\n",
        "  all_clf_labels_curr = pickle.load(f)[superclass_idx]\n",
        "\n",
        "reorder_idx = []\n",
        "for subclass in train_subclasses[superclass_idx]:\n",
        "  reorder_idx.extend(np.where(subclass_labels == subclass)[0])\n",
        "\n",
        "fig = plt.figure(figsize=(20, 10))\n",
        "for plot_idx, all_clf_labels in enumerate([all_clf_labels_prev, all_clf_labels_curr]):\n",
        "  all_clf_labels = all_clf_labels[reorder_idx]\n",
        "  cluster_assignment = np.zeros((n_data, n_data), dtype=int)\n",
        "  for cluster_idx in list(set(all_clf_labels)):\n",
        "    instance_idx = np.where(all_clf_labels == cluster_idx)[0]\n",
        "    from itertools import combinations\n",
        "    for pair in list(combinations(instance_idx, 2)):\n",
        "      cluster_assignment[pair[0]][pair[1]] = 1\n",
        "      cluster_assignment[pair[1]][pair[0]] = 1\n",
        "\n",
        "  ax = fig.add_subplot(1, 2, plot_idx+1)\n",
        "  #fig, ax = plt.subplots(1,1, figsize=(15,15))\n",
        "  img = ax.matshow(cluster_assignment)\n",
        "  n_subclass_per_class = len(train_subclasses[superclass_idx])\n",
        "  subclass_names = [hier.LEAF_NUM_TO_NAME[s].split(',')[0] for s in train_subclasses[superclass_idx]]\n",
        "  ax.set_xticks(range(0, n_data, n_data//n_subclass_per_class))\n",
        "  ax.set_xticklabels(subclass_names, rotation=90)\n",
        "  ax.set_yticks(range(0, n_data, n_data//n_subclass_per_class))\n",
        "  ax.set_yticklabels(subclass_names)\n",
        "  ax.grid(color='w', linewidth=1)\n",
        "  #fig.colorbar(img)\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MOWnxMdQ5JOV"
      },
      "outputs": [],
      "source": [
        "superclass_idx = 0\n",
        "print(\"----------SUPERCLASS:\", label_map[superclass_idx])\n",
        "subclass_idx = np.array([i for i in range(len(all_subclass_labels)) if all_subclass_labels[i] in train_subclasses[superclass_idx]])\n",
        "subclass_labels = all_subclass_labels[subclass_idx]\n",
        "n_data = len(subclass_idx)\n",
        "n_subclass_per_class = len(train_subclasses[superclass_idx])\n",
        "n_data_per_subclass = n_data // n_subclass_per_class\n",
        "\n",
        "model_path = os.path.join(BASE_DIR, 'breeds/breeds_training_level_3_nsubclass_20_400_epochs_ema_0.99/')\n",
        "reorder_idx = []\n",
        "for subclass in train_subclasses[superclass_idx]:\n",
        "  reorder_idx.extend(np.where(subclass_labels == subclass)[0])\n",
        "\n",
        "fig = plt.figure(figsize=(20, 10))\n",
        "plot_data = [[] for _ in range(n_subclass_per_class)]\n",
        "for ckpt_number in range(11, 162, 10):\n",
        "  with gfile.Open(os.path.join(model_path, f'clf_labels_ckpt_{ckpt_number}.pkl'), 'rb') as f:\n",
        "    all_clf_labels = pickle.load(f)[superclass_idx][reorder_idx]\n",
        "\n",
        "  cluster_assignment = np.zeros((n_data, n_data), dtype=int)\n",
        "  for cluster_idx in list(set(all_clf_labels)):\n",
        "    instance_idx = np.where(all_clf_labels == cluster_idx)[0]\n",
        "    from itertools import combinations\n",
        "    for pair in list(combinations(instance_idx, 2)):\n",
        "      cluster_assignment[pair[0]][pair[1]] = 1\n",
        "      cluster_assignment[pair[1]][pair[0]] = 1\n",
        "  for i, _ in enumerate(train_subclasses[superclass_idx]):\n",
        "    count = np.sum(cluster_assignment[n_data_per_subclass*i : (n_data_per_subclass*(i+1)), n_data_per_subclass*i : (n_data_per_subclass*(i+1))]) / n_data_per_subclass**2\n",
        "    plot_data[i].append(count)\n",
        "\n",
        "for i, s in enumerate(train_subclasses[superclass_idx]):\n",
        "  plt.plot([(c-1) * 1000 for c in range(11, 162, 10)], plot_data[i], label=hier.LEAF_NUM_TO_NAME[s])\n",
        "plt.legend()\n",
        "plt.xlabel(\"Step\")\n",
        "plt.ylabel('Recall')\n",
        "plt.show()"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "last_runtime": {
        "build_target": "//learning/deepmind/dm_python:dm_notebook3",
        "kind": "private"
      },
      "private_outputs": true,
      "provenance": [
        {
          "file_id": "1Twuiw6Yshv0fHVYHnpXT5VfWTHCNawbc",
          "timestamp": 1664586424113
        }
      ]
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
